package ai.koog.agents.core.agent

import ai.koog.agents.core.agent.GraphAIAgent.FeatureContext
import ai.koog.agents.core.agent.config.AIAgentConfig
import ai.koog.agents.core.agent.entity.AIAgentGraphStrategy
import ai.koog.agents.core.annotation.InternalAgentsApi
import ai.koog.agents.core.tools.Tool
import ai.koog.agents.core.tools.ToolRegistry
import ai.koog.agents.core.tools.annotations.InternalAgentToolsApi
import ai.koog.prompt.dsl.prompt
import ai.koog.prompt.executor.model.PromptExecutor
import ai.koog.prompt.llm.LLModel
import ai.koog.prompt.params.LLMParams
import ai.koog.prompt.processor.ResponseProcessor
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import kotlinx.datetime.Clock
import kotlinx.serialization.KSerializer
import kotlinx.serialization.serializer
import kotlin.reflect.KType
import kotlin.reflect.typeOf

/**
 * [AIAgentService] is a core interface for managing AI agents. The service allows creation, removal, and
 * management of AI agents and provides functionalities to list agents based on their statuses.f
 *
 * A single instance of [AIAgentService] manages one kind of uniform AI Agents serving same purpose and solving the same type
 * of user task. It's useful to create, manage, and track the progress of running agents solving similar user tasks in parallel.
 *
 * @param Input Type parameter for the input data for the agents.
 * @param Output Type parameter for the output data generated by the agents.
 */
public interface AIAgentService<Input, Output, TAgent : AIAgent<Input, Output>> {

    /**
     * Represents an instance of PromptExecutor responsible for executing prompt-based operations.
     *
     * This variable is used to handle interactions and executions related to prompts,
     * providing a centralized mechanism to manage user inputs or automated instructions.
     */
    public val promptExecutor: PromptExecutor

    /**
     * Holds the configuration details for an AI agent.
     *
     * This property encapsulates the settings and parameters
     * required to initialize and operate an AI agent, such as
     * behavior configurations, thresholds, or any specific operational
     * constraints. It is essential for defining the agent's behavior and performance.
     */
    public val agentConfig: AIAgentConfig

    /**
     * Represents the central repository or registry for managing tools within the system.
     *
     * This variable provides access to a tool registry that can be used to store, retrieve,
     * and manage different tools registered in the application.
     *
     * It is expected to be immutable and thread-safe, ensuring consistent behavior across the application.
     */
    public val toolRegistry: ToolRegistry

    /**
     * Removes the specified AI agent from the service.
     *
     * @param agent The AI agent to be removed.
     * @return True if the agent was successfully removed, or false if the agent was not found.
     */
    public suspend fun removeAgent(agent: TAgent): Boolean

    /**
     * Removes an AI agent based on its unique identifier.
     *
     * @param id The unique identifier of the AI agent to be removed.
     * @return True if the agent with the specified ID was successfully removed, or false if no such agent was found.
     */
    public suspend fun removeAgentWithId(id: String): Boolean

    /**
     * Retrieves an AI agent based on its unique identifier.
     *
     * @param id The unique identifier of the AI agent to retrieve.
     * @return The AI agent associated with the specified ID, or null if no agent is found.
     */
    public suspend fun agentById(id: String): TAgent?

    /**
     * Retrieves a comprehensive list of all AI agents currently managed by the service,
     * regardless of their state (active, inactive, or finished).
     *
     * @return A list of all AI agents managed by the service.
     */
    public suspend fun listAllAgents(): List<TAgent>

    /**
     * Retrieves a list of active AI agents currently managed by the service.
     *
     * @return A list containing the currently active AI agents.
     */
    public suspend fun listActiveAgents(): List<TAgent>

    /**
     * Retrieves a list of inactive AI agents currently managed by the service.
     *
     * @return A list of AI agents that are marked as inactive.
     */
    public suspend fun listInactiveAgents(): List<TAgent>

    /**
     * Retrieves a list of AI agents that have completed their tasks and are marked as finished.
     *
     * @return A list of finished AI agents.
     */
    public suspend fun listFinishedAgents(): List<TAgent>

    /**
     * Closes all AI agents currently managed by the service.
     *
     * This method retrieves the list of all agents, regardless of their state (active, inactive, or finished),
     * and invokes the [AIAgent.close] function on each of them, releasing any underlying resources.
     */
    public suspend fun closeAll() {
        listAllAgents().forEach { it.close() }
    }

    /**
     * Companion object that provides factory methods for creating instances of
     * GraphAIAgentService with various configurations.
     */
    // TODO move out and replace invoke methods with proper factories
    public companion object {
        /**
         * Converts a given [GraphAIAgent] instance into an [AIAgentService] instance.
         *
         * @param Input The input type that the agent processes.
         * @param Output The output type that the agent produces.
         * @param agent The [GraphAIAgent] to be converted into a service instance.
         * @return An [AIAgentService] instance constructed from the provided [GraphAIAgent].
         */
        @OptIn(InternalAgentsApi::class)
        public inline fun <reified Input, reified Output> fromAgent(
            agent: GraphAIAgent<Input, Output>
        ): AIAgentService<Input, Output, GraphAIAgent<Input, Output>> =
            AIAgentService(
                promptExecutor = agent.promptExecutor,
                agentConfig = agent.agentConfig,
                strategy = agent.strategy,
                toolRegistry = agent.toolRegistry,
                installFeatures = agent.installFeatures
            )

        /**
         * Creates a new instance of [AIAgentService] by transforming a given [FunctionalAIAgent].
         *
         * @param Input The type of input data expected by the agent.
         * @param Output The type of output data produced by the agent.
         * @param agent The [FunctionalAIAgent] to be transformed into an [AIAgentService].
         * @return A new [AIAgentService] instance configured with the parameters of the provided agent.
         */
        @OptIn(InternalAgentsApi::class)
        public fun <Input, Output> fromAgent(
            agent: FunctionalAIAgent<Input, Output>
        ): AIAgentService<Input, Output, FunctionalAIAgent<Input, Output>> =
            AIAgentService(
                promptExecutor = agent.promptExecutor,
                agentConfig = agent.agentConfig,
                strategy = agent.strategy,
                toolRegistry = agent.toolRegistry,
                installFeatures = agent.installFeatures
            )

        /**
         * Invokes the creation of a [GraphAIAgentService] instance with the provided configuration, strategy,
         * tool registry, and optional feature installation logic.
         *
         * @param promptExecutor The executor responsible for processing AI prompts and responses.
         * @param agentConfig Configuration parameters for the AI agent.
         * @param strategy A strategy defining the graph structure for AI agent interactions and processing.
         * @param toolRegistry*/
        @OptIn(InternalAgentsApi::class)
        public inline operator fun <reified Input, reified Output> invoke(
            promptExecutor: PromptExecutor,
            agentConfig: AIAgentConfig,
            strategy: AIAgentGraphStrategy<Input, Output>,
            toolRegistry: ToolRegistry = ToolRegistry.EMPTY,
            noinline installFeatures: FeatureContext.() -> Unit = {},
        ): GraphAIAgentService<Input, Output> = GraphAIAgentService(
            promptExecutor = promptExecutor,
            agentConfig = agentConfig,
            strategy = strategy,
            inputType = typeOf<Input>(),
            outputType = typeOf<Output>(),
            toolRegistry = toolRegistry,
            installFeatures = installFeatures
        )

        /**
         * Invokes the creation of a [GraphAIAgentService] with the provided dependencies, configuration,
         * and optional parameters for customization.
         *
         * @param promptExecutor The executor responsible for handling prompt-based interactions.
         * @param llmModel The large language model to be used by the agent.
         * @param strategy The graph strategy defining the agent's execution behavior. Defaults to a single-run*/
        public operator fun invoke(
            promptExecutor: PromptExecutor,
            llmModel: LLModel,
            responseProcessor: ResponseProcessor? = null,
            strategy: AIAgentGraphStrategy<String, String> = singleRunStrategy(),
            toolRegistry: ToolRegistry = ToolRegistry.EMPTY,
            systemPrompt: String = "",
            temperature: Double = 1.0,
            numberOfChoices: Int = 1,
            maxIterations: Int = 50,
            installFeatures: FeatureContext.() -> Unit = {}
        ): GraphAIAgentService<String, String> = AIAgentService(
            promptExecutor = promptExecutor,
            strategy = strategy,
            agentConfig = AIAgentConfig(
                prompt = prompt(
                    id = "chat",
                    params = LLMParams(
                        temperature = temperature,
                        numberOfChoices = numberOfChoices
                    )
                ) {
                    system(systemPrompt)
                },
                model = llmModel,
                maxAgentIterations = maxIterations,
                responseProcessor = responseProcessor
            ),
            toolRegistry = toolRegistry,
            installFeatures = installFeatures
        )

        /**
         * Invokes the creation of a FunctionalAIAgentService instance with the provided parameters.
         *
         * @param promptExecutor The executor responsible for handling prompts and managing their execution.
         * @param agentConfig The configuration parameters for the AI agent.
         * @param strategy The functional strategy that defines the behavior and capabilities of the AI agent.
         * @param toolRegistry The registry containing tools that can be used by the agent. Defaults to an empty registry if not specified.
         * @param installFeatures A lambda expression to configure and install additional features to the AI agent context.
         * @return An instance of FunctionalAIAgentService initialized with the given parameters.
         */
        @OptIn(InternalAgentsApi::class)
        public operator fun <Input, Output> invoke(
            promptExecutor: PromptExecutor,
            agentConfig: AIAgentConfig,
            strategy: AIAgentFunctionalStrategy<Input, Output>,
            toolRegistry: ToolRegistry = ToolRegistry.EMPTY,
            installFeatures: FunctionalAIAgent.FeatureContext.() -> Unit = {},
        ): FunctionalAIAgentService<Input, Output> = FunctionalAIAgentService(
            promptExecutor = promptExecutor,
            agentConfig = agentConfig,
            toolRegistry = toolRegistry,
            strategy = strategy,
            installFeatures = installFeatures
        )
    }

    /**
     * Creates a new AI agent with the specified configuration, tool registry, and optional parameters.
     *
     * @param id An optional unique identifier for the agent. If null, a default identifier will be generated.
     * @param additionalToolRegistry The registry of tools to be used by the agent. Defaults to an empty registry.
     * @param agentConfig The configuration settings for the AI agent. Defaults to the service's agent configuration.
     * @param clock The clock instance to manage time-related operations, defaulting to the system clock.
     * @return The newly created AI agent.
     */
    public suspend fun createAgent(
        id: String? = null,
        additionalToolRegistry: ToolRegistry = ToolRegistry.EMPTY,
        agentConfig: AIAgentConfig = this.agentConfig,
        clock: Clock = Clock.System,
    ): TAgent

    /**
     * Creates an AI agent with the specified configuration and runs it using the provided input.
     *
     * @param agentInput The input data to be processed by the agent.
     * @param id An optional unique identifier for the agent. If null, a default identifier will be generated.
     * @param additionalToolRegistry The tool registry to be used by the agent. Defaults to the service's tool registry.
     * @param agentConfig The configuration parameters for the AI agent. Defaults to the service's agent configuration.
     * @param clock The clock instance used for managing time-related operations. Defaults to the system clock.
     * @return The output generated by the agent after processing the input.
     */
    public suspend fun createAgentAndRun(
        agentInput: Input,
        id: String? = null,
        additionalToolRegistry: ToolRegistry = this.toolRegistry,
        agentConfig: AIAgentConfig = this.agentConfig,
        clock: Clock = Clock.System,
    ): Output = createAgent(id, additionalToolRegistry, agentConfig, clock).run(agentInput)
}

/**
 * Abstract base class for implementing an AI Agent Service. This class provides
 * foundational functionality for managing and interacting with AI agents.
 *
 * The class maintains a list of AI agents under its management and defines
 * common operations, such as creation, removal, querying, and shutting down agents.
 * Concrete implementations are expected to define the behavior for creating managed agents.
 *
 * @param Input the input type expected by the AI agents.
 * @param Output the output type expected from the AI agents.
 */
public abstract class AIAgentServiceBase<Input, Output, TAgent : AIAgent<Input, Output>> :
    AIAgentService<Input, Output, TAgent> {
    /**
     * A mutable map that holds the agents currently managed by the AIAgentServiceBase instance.
     *
     * This list is used to track all agents created and managed by the service, allowing for operations
     * such as addition, removal, and querying of agents based on their state or lifecycle.
     */
    private val managedAgents: MutableMap<String, TAgent> = mutableMapOf()
    private val managedAgentsMutex = Mutex()

    /**
     * Creates and registers a new managed AI agent with the specified configuration and tool registry.
     *
     * @param id An optional unique identifier for the AI agent. If null, a default identifier will be generated.
     * @param additionalToolRegistry A tool registry with additional tools available to the AI agent.
     * @param agentConfig The configuration for the AI agent, including settings for its behavior and capabilities.
     * @param clock The clock instance used for managing time-related operations. Defaults to the system clock.
     * @return A new instance of `AIAgent` initialized with the specified parameters.
     */
    @InternalAgentsApi
    public abstract fun createManagedAgent(
        id: String? = null,
        additionalToolRegistry: ToolRegistry,
        agentConfig: AIAgentConfig,
        clock: Clock = Clock.System,
    ): TAgent

    /**
     * Creates and registers a new managed AI agent with the specified configuration and tool registry.
     *
     * @param id An optional unique identifier for the AI agent. If null, a default identifier will be generated.
     * @param additionalToolRegistry A tool registry with additional tools available to the AI agent.
     * @param agentConfig The configuration for the AI agent, including settings for its behavior and capabilities.
     * @param clock The clock instance used for managing time-related operations.
     * @return A new instance of `AIAgent` initialized with the specified parameters.
     */
    @OptIn(InternalAgentsApi::class)
    final override suspend fun createAgent(
        id: String?,
        additionalToolRegistry: ToolRegistry,
        agentConfig: AIAgentConfig,
        clock: Clock
    ): TAgent = managedAgentsMutex.withLock {
        val agent = createManagedAgent(id, additionalToolRegistry, agentConfig, clock)
        managedAgents[agent.id] = agent
        return agent
    }

    /**
     * Removes the specified AI agent from the managed collection.
     *
     * @param agent The AI agent to be removed. This agent must be part of the managed agents' collection.
     * @return `true` if the agent was successfully removed, or `false` if the agent was not found or could not be removed.
     */
    final override suspend fun removeAgent(agent: TAgent): Boolean = removeAgentWithId(agent.id)

    /**
     * Removes a managed AI agent identified by the specified unique identifier.
     *
     * @param id The unique identifier of the AI agent to be removed.
     * @return `true` if the agent was successfully removed, `false` if no agent with the given ID exists.
     */
    final override suspend fun removeAgentWithId(id: String): Boolean = managedAgentsMutex.withLock {
        managedAgents.remove(id) != null
    }

    /**
     * Retrieves an AI agent managed by this service based on the provided unique identifier.
     *
     * @param id The unique identifier of the AI agent to be retrieved.
     * @return The AI agent corresponding to the given identifier, or `null` if no agent with the specified ID exists.
     */
    final override suspend fun agentById(id: String): TAgent? = managedAgentsMutex.withLock {
        managedAgents[id]
    }

    /**
     * Retrieves a list of all currently active AI agents managed by the service.
     *
     *
     * @return A list of active AI agents of type `AIAgent<Input, Output>`.
     */
    final override suspend fun listActiveAgents(): List<TAgent> = managedAgentsMutex.withLock {
        managedAgents.filterValues { it.isRunning() }.values.toList()
    }

    /**
     * Retrieves a list of inactive AI agents managed by this service.
     *
     * An agent is considered inactive if it is not currently running.
     *
     * @return A list of AI agents that are inactive.
     */
    final override suspend fun listInactiveAgents(): List<TAgent> = managedAgentsMutex.withLock {
        managedAgents.filterValues { !it.isRunning() }.values.toList()
    }

    /**
     * Retrieves a list of all agents that have already finished their task.
     *
     * @return A list of finished agents managed by this service.
     */
    final override suspend fun listFinishedAgents(): List<TAgent> = managedAgentsMutex.withLock {
        managedAgents.filterValues { it.isFinished() }.values.toList()
    }

    /**
     * Retrieves a list of all AI agents currently managed by the service, regardless of their state
     * (e.g., active, inactive, or finished).
     *
     * @return A list of all AI agents of type `AIAgent<Input, Output>` managed by this service.
     */
    override suspend fun listAllAgents(): List<TAgent> = managedAgentsMutex.withLock {
        managedAgents.values.toList()
    }
}

/**
 * A service class for managing graph-based AI agents, responsible for orchestrating the behavior
 * and execution of agents within a configurable graph-based strategy.
 *
 * This class extends [AIAgentServiceBase] and provides an implementation for creating
 * AI agents that follow a graph-based approach. It leverages a strategy pattern
 * defined by [AIAgentGraphStrategy] to handle input and output transformations.
 *
 * The [GraphAIAgentService] is marked with the [InternalAgentsApi] annotation,
 * indicating that it is intended for internal use within agent-related implementations and not
 * guaranteed for backward compatibility.
 *
 * @param promptExecutor An instance of [PromptExecutor], responsible for executing prompts
 *        using language models and handling responses.
 **/
@Suppress("MissingKDocForPublicAPI")
public class GraphAIAgentService<Input, Output>
@InternalAgentsApi
public constructor(
    override val promptExecutor: PromptExecutor,
    override val agentConfig: AIAgentConfig,
    public val strategy: AIAgentGraphStrategy<Input, Output>,
    private val inputType: KType,
    private val outputType: KType,
    override val toolRegistry: ToolRegistry,
    public val installFeatures: FeatureContext.() -> Unit
) : AIAgentServiceBase<Input, Output, GraphAIAgent<Input, Output>>() {

    @InternalAgentsApi
    override fun createManagedAgent(
        id: String?,
        additionalToolRegistry: ToolRegistry,
        agentConfig: AIAgentConfig,
        clock: Clock,
    ): GraphAIAgent<Input, Output> = GraphAIAgent(
        inputType = inputType,
        outputType = outputType,
        promptExecutor = promptExecutor,
        agentConfig = agentConfig,
        toolRegistry = toolRegistry + additionalToolRegistry,
        strategy = strategy,
        id = id,
        clock = clock,
        installFeatures = installFeatures
    )
}

/**
 * A service for managing functional AI agents that operate based on a specified execution strategy.
 *
 * This class provides the infrastructure for creating and managing functional AI agents.
 * It integrates an execution strategy with associated input and output types, a tool registry for
 * additional functionalities, and a customizable feature installation context.
 *
 * @param Input The type of input data that the functional AI agent will process.
 * @param Output The type of output data produced by the functional AI agent.
 * @property promptExecutor The executor responsible for handling and executing prompts.
 * @property agentConfig The configuration settings for the AI agent, including model and prompt behavior.
 * @property strategy The functional strategy defining the agent's iteration logic and processing behavior.
 */
@Suppress("MissingKDocForPublicAPI")
public class FunctionalAIAgentService<Input, Output>
@InternalAgentsApi
public constructor(
    override val promptExecutor: PromptExecutor,
    override val agentConfig: AIAgentConfig,
    public val strategy: AIAgentFunctionalStrategy<Input, Output>,
    override val toolRegistry: ToolRegistry,
    public val installFeatures: FunctionalAIAgent.FeatureContext.() -> Unit
) : AIAgentServiceBase<Input, Output, FunctionalAIAgent<Input, Output>>() {

    /**
     * Creates and returns a managed instance of [FunctionalAIAgent].
     *
     * @param clock The clock instance used for time-related operations within the agent.
     * @param id The identifier for the agent. Random UUID will be generated if set to null.
     * @return A managed AI agent instance implementing the AIAgent interface.
     */
    @InternalAgentsApi
    override fun createManagedAgent(
        id: String?,
        additionalToolRegistry: ToolRegistry,
        agentConfig: AIAgentConfig,
        clock: Clock,
    ): FunctionalAIAgent<Input, Output> = FunctionalAIAgent(
        promptExecutor = promptExecutor,
        agentConfig = agentConfig,
        toolRegistry = toolRegistry + additionalToolRegistry,
        strategy = strategy,
        id = id,
        clock = clock,
        installFeatures = installFeatures
    )
}

/**
 * Invokes the process to create and return an instance of `GraphAIAgentService`.
 *
 * @param promptExecutor The executor responsible for handling prompts during the agent's operation.
 * @param agentConfig The configuration object for the AI agent.
 * @param strategy The strategy defining how the agent processes tasks and connections within the graph. Defaults to `singleRunStrategy`.
 * @param toolRegistry The registry containing tools available for the agent to use. Defaults to an empty*/
@OptIn(InternalAgentsApi::class)
public operator fun AIAgentService.Companion.invoke(
    promptExecutor: PromptExecutor,
    agentConfig: AIAgentConfig,
    strategy: AIAgentGraphStrategy<String, String> = singleRunStrategy(),
    toolRegistry: ToolRegistry = ToolRegistry.EMPTY,
    installFeatures: FeatureContext.() -> Unit = {},
): GraphAIAgentService<String, String> = GraphAIAgentService(
    promptExecutor = promptExecutor,
    agentConfig = agentConfig,
    strategy = strategy,
    inputType = typeOf<String>(),
    outputType = typeOf<String>(),
    toolRegistry = toolRegistry,
    installFeatures = installFeatures
)

/**
 * Creates an [AIAgent] and converts it to a [Tool] that can be used by other AI Agents.
 *
 * @param agentName Agent name that would be a tool name for this agent tool.
 * @param agentDescription Agent description that would be a tool description for this agent tool.
 * @param inputDescription An optional description of the agent's input. Required for primitive types only!
 *  * If not specified for a primitive input type (ex: String, Int, ...), an empty input description will be sent to LLM.
 *  * Does not have any effect for non-primitive [Input] type with @LLMDescription annotations.
 * @param inputSerializer Serializer to deserialize tool arguments to agent input.
 * @param outputSerializer Serializer to serialize agent output to a tool result.
 * @return A special tool that wraps the agent functionality.
 * @param parentAgentId Optional ID of the parent AI agent. Tool agent IDs will be generated as "parentAgentId.<number of tool call>"
 * @param clock The clock instance used to manage time-related operations. Defaults to `Clock.System`.
 * @return A tool instance configured with the provided parameters, representing the AI agent.
 */
@OptIn(InternalAgentToolsApi::class)
public inline fun <reified Input, reified Output> AIAgentService<Input, Output, *>.createAgentTool(
    agentName: String,
    agentDescription: String,
    inputDescription: String? = null,
    inputSerializer: KSerializer<Input> = serializer(),
    outputSerializer: KSerializer<Output> = serializer(),
    parentAgentId: String? = null,
    clock: Clock = Clock.System
): Tool<Input, AIAgentTool.AgentToolResult<Output>> = AIAgentTool(
    agentService = this,
    agentName = agentName,
    agentDescription = agentDescription,
    inputDescription = inputDescription,
    inputSerializer = inputSerializer,
    outputSerializer = outputSerializer,
    parentAgentId = parentAgentId
)
